package rate

import (
	"context"

	"go-toolkit/library/ecode"
	"go-toolkit/library/log"
	"go-toolkit/library/stat/metric"

	"golang.org/x/time/rate"
	"google.golang.org/grpc"
)

const (
	_defBurst = 100
)

var (
	_metricServerTB = metric.NewCounterVec(&metric.CounterVecOpts{
		Namespace: "grpc_server",
		Subsystem: "",
		Name:      "tokenbucket_total",
		Help:      "grpc server tokenbucket total.",
		Labels:    []string{"rulename", "rulekey"},
	})
)

// Limiter controls how frequently events are allowed to happen.
type Limiter struct {
	rbs map[string]*ruleBucket
}

// Config limitter  conf.
type Config []*Limit

// Limit limit conf.
type Limit struct {
	Rule  string
	Key   string
	Limit rate.Limit
	Burst int
}

// New return Limiter.
func New(conf Config, rs ...Rule) (l *Limiter) {
	// 默认内置 rule
	rs = append(rs, RuleCaller{}, RuleMethod{}, RuleMethodCaller{})

	rbs := make(map[string]*ruleBucket, len(rs))
	for _, r := range rs {
		if _, ok := rbs[r.Name()]; ok {
			continue
		}

		rb := newRuleBucket(r)
		rbs[r.Name()] = rb
	}

	l = &Limiter{rbs}
	l.Reload(conf)
	return
}

// Reload reload limit conf.
func (l *Limiter) Reload(conf Config) {
	m := make(map[string]Config)

	for _, c := range conf {
		if c.Rule == "" || c.Key == "" {
			continue
		}

		m[c.Rule] = append(m[c.Rule], c)
	}

	for _, v := range l.rbs {
		v.Reload(m[v.r.Name()])
	}
}

func (l *Limit) fix() (lim rate.Limit, b int) {
	lim = rate.Inf
	b = _defBurst
	if l.Limit <= 0 {
		lim = rate.Inf
	} else {
		lim = l.Limit
	}
	if l.Burst > 0 {
		b = l.Burst
	}
	return
}

// Allow reports whether event may happen at time now.
func (l *Limiter) Allow(ctx context.Context, req interface{}, args *grpc.UnaryServerInfo) (ok bool, name string, key string) {
	ok = true
	for k, v := range l.rbs {
		ok, key = v.Allow(ctx, req, args)
		if !ok {
			name = k
			return
		}
	}
	return
}

// Limit is a server interceptor that detects and rejects overloaded traffic.
func (l *Limiter) Limit() grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req interface{}, args *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
		ok, name, key := l.Allow(ctx, req, args)
		if !ok {
			_metricServerTB.Inc(name, key)
			err = ecode.LimitExceed
			log.Warn("warden tokenbucket limit name(%s) key(%s)", name, key)
			return
		}
		resp, err = handler(ctx, req)
		return
	}
}
